#!/usr/bin/env python
# coding=utf-8
import akshare as ak
import pandas as pd
import numpy as np
from include import dataBase

class dealCode:
    def __init__(self,connect):
        self.connect = connect
        self.baseName = "codeX"


    def save(self):
        # 获取上海所有的股票代码
        dataSH = ak.stock_sh_a_spot_em()
        codeSH = dataSH[['代码','名称']]

        # 获取深圳所有的股票代码
        dataSZ = ak.stock_sz_a_spot_em()
        codeSZ = dataSZ[['代码','名称']]

        # 拼接得到所有的股票代码
        codeAll = pd.concat([codeSH,codeSZ]).reset_index(drop=True)

        sql = ''' CREATE TABLE IF NOT EXISTS {} (
        code VARCHAR(10) PRIMARY KEY,
        name VARCHAR(100) NOT NULL,
        market VARCHAR(10) NOT NULL
        ); '''.format(self.baseName)
        self.connect.create(self.baseName,sql)

        for date in codeAll.itertuples():
            market = self.getMarket(date[1])
            sql = '''insert into {} values("{}","{}","{}")ON DUPLICATE KEY UPDATE code = code;'''.format(self.baseName,date[1],date[2],market)
            self.connect.insert(self.baseName,sql)
        print("更新保存所有代码完成")

    def getMarket(self,code):
        ''' 根据代码返回所在市场 '''
        if code[0:3] in ["688","689"]:
            market = "科创板"
        elif code[0] == "3":
            market = "创业板"
        elif code[0:3] in ['600','601','603','605']:
            market = "上海主板"
        elif code[0] == '0':
            market = "深圳主板"
        elif code[0] in ['4','8','9']:
            market = "北交所"
        else:
            market = "新类型"
            print(code,"新的代码类型")

        return market

    def getAllCode(self):
        sql = "select code from {};".format(self.baseName)
        return self.connect.select(self.baseName,sql)

    def getMarketCode(self,market):
        sql = "select code from {} where market = '{}';".format(self.baseName,market)
        return self.connect.select(self.baseName,sql)

    def getName(self,code):
        sql = "select name from {} where code = '{}';".format(self.baseName,code)
        return self.connect.select(self.baseName,sql)

    def getDateCode(self,date):
        pass









